"""
加载预训练好的模型并测试效果。
"""

import time  # 导入时间模块，用于计算推理耗时
from typing import List  # 导入 List 类型提示，用于类型标注
import torch  # 导入 PyTorch 模块，用于深度学习推理
from rich import print  # 导入 rich 库，用于美化打印输出
from transformers import AutoTokenizer, AutoModelForMaskedLM  # 导入 Hugging Face 的分词器和模型
# import sys
from utils.verbalizer import Verbalizer  # 导入 Verbalizer 类，用于标签映射
from data_handle.template import HardTemplate  # 导入 HardTemplate 类，用于模板处理
from data_handle.data_preprocess import convert_example  # 导入数据预处理函数
from utils.common_utils import convert_logits_to_ids  # 导入 logits 转换为 token IDs 的工具函数

device = 'cpu:0'  # 设置设备为 CPU
# device = 'cuda:0'  # 如果使用 GPU，取消注释此行（适用于 Windows 环境）
model_path = 'checkpoints/model_best_old'  # 指定预训练模型路径
tokenizer = AutoTokenizer.from_pretrained(model_path)  # 加载预训练分词器
model = AutoModelForMaskedLM.from_pretrained(model_path)  # 加载预训练模型
model.to(device).eval()  # 将模型移动到指定设备并切换到评估模式

max_label_len = 2  # 标签最大长度
verbalizer = Verbalizer(
    verbalizer_file='data/verbalizer.txt',  # 加载 verbalizer 文件
    tokenizer=tokenizer,  # 使用分词器
    max_label_len=max_label_len  # 设置最大标签长度
)
prompt = open('data/prompt.txt', 'r', encoding='utf8').readlines()[0].strip()  # 读取 prompt 内容
hard_template = HardTemplate(prompt=prompt)  # 初始化模板转换器
print(f'Prompt is -> {prompt}')  # 打印使用的 prompt


def inference(contents: List[str]):
    """
    推理函数，输入原始句子，输出 mask label 的预测值。

    Args:
        contents (List[str]): 原始句子列表。
    """
    with torch.no_grad():  # 禁用梯度计算，节省内存并加速推理
        start_time = time.time()  # 记录推理开始时间
        examples = {'text': contents}  # 构造输入数据字典
        tokenized_output = convert_example(
            examples,  # 输入数据
            tokenizer,  # 分词器
            hard_template=hard_template,  # 模板转换器
            max_seq_len=128,  # 最大序列长度
            max_label_len=max_label_len,  # 最大标签长度
            train_mode=False,  # 非训练模式
            return_tensor=True  # 返回张量格式
        )  # 数据预处理
        logits = model(
            input_ids=tokenized_output['input_ids'].to(device),  # 输入 ID
            token_type_ids=tokenized_output['token_type_ids'].to(device),  # 句子类型 ID
            attention_mask=tokenized_output['attention_mask'].to(device)  # 注意力掩码
        ).logits  # 获取模型输出的 logits
        predictions = convert_logits_to_ids(logits, tokenized_output['mask_positions']).cpu().numpy().tolist()  # 提取 mask 位置的预测结果
        predictions = verbalizer.batch_find_main_label(predictions)  # 找到子标签对应的主标签
        predictions = [ele['label'] for ele in predictions]  # 提取主标签名称
        used = time.time() - start_time  # 计算推理耗时
        print(f'Used {used}s.')  # 打印耗时信息
        return predictions  # 返回预测结果


if __name__ == '__main__':
    contents = [
        '这个地方的天台景色迷人，躺在舒适的躺椅上，感觉无比放松。由于这里组织的一些活动，我认为这里的性价比非常高，特别适合家庭旅行。尤其是前往迪士尼游玩时，这里的地理位置非常便利。我已经决定下次有机会还会再来，并强烈推荐给我的朋友们。',
        '这家酒店在环境、设施等方面都给人留下了深刻的印象。周围的生活配套设施一应俱全，而前台小姐姐特别美好喜欢。早餐的质量也很高，整体来说，这是一家性价比极高的酒店，无论是硬件设施还是软件服务都非常出色，值得向所有人推荐。',
        "没想到物流这么快，商品第二天就到了。虽然还没有使用，但考虑到它的便捷性和小巧的设计，对于外出旅行来说简直是完美选择。这次愉快的购物体验让我对未来的购买充满信心。",
        "因为喜欢的一家面包店参加了福行市的无早集市，所以特地前来参观。在这里找到了一直喜爱的巧克力磅蛋糕和其他美味的小点心。虽然有些遗憾柠檬磅蛋糕暂时缺货，但我对未来的新口味充满了期待。这里的饼干类甜点同样让人爱不释手。",
        "这次选择了上海的一家酒店，房间布置温馨舒适，特别是床铺的质量非常好，晚上睡得很安稳。隔音效果也不错，确保了夜晚的宁静。孩子们也非常喜欢这里的氛围，这一切都让我们决定下次再来上海时，依然选择这家酒店。"
    ]  # 输入评论内容
    print("针对下面的文本评论，请分别给出对应所属类别：")  # 提示信息
    res = inference(contents)  # 调用推理函数

    # print('inference label(s):', res)  # 打印推理结果
    new_dict = {}  # 初始化结果字典
    for i in range(len(contents)):  # 遍历评论内容
        new_dict[contents[i]] = res[i]  # 将评论与预测类别存入字典
    print(new_dict)  # 打印最终结果